#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Cookie 管理工具
提供获取和保存 Cookie 的功能
"""

import asyncio
import json
import time
from pathlib import Path
from playwright.async_api import async_playwright


async def get_and_save_cookies(
    url: str = "https://www.douyin.com",
    cookie_file: str = "douyin_storage_state.json",
    wait_time: int = 60,
    extend_expiry_days: int = 365
):
    """
    打开浏览器，等待用户手动登录，然后保存 Cookie（storage_state 格式）
    
    参数:
        url: 要打开的网址
        cookie_file: Cookie 保存的文件名（storage_state 格式）
        wait_time: 等待用户登录的时间（秒）
        extend_expiry_days: 延长 cookie 过期时间的天数（默认 365 天）
    """
    browser = None
    context = None
    
    try:
        print("=" * 60)
        print("Cookie 获取工具")
        print("=" * 60)
        print(f"\n📋 目标网址: {url}")
        print(f"💾 保存位置: {Path(cookie_file).absolute()}")
        print(f"⏱️  等待时间: {wait_time} 秒")
        print("\n" + "=" * 60)
        
        # 使用 Playwright 启动浏览器
        print("\n正在启动浏览器...")
        playwright = await async_playwright().start()
        browser = await playwright.chromium.launch(
            headless=False,
            channel='chrome',
            args=[
                '--disable-blink-features=AutomationControlled',
                '--no-sandbox',
                '--disable-dev-shm-usage',
            ]
        )
        
        # 创建浏览器上下文
        context = await browser.new_context(
            viewport={'width': 1920, 'height': 1080},
            user_agent='Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36'
        )
        
        print("✅ 浏览器已启动")
        
        # 打开目标网址
        print(f"\n正在打开网址: {url}")
        page = await context.new_page()
        await page.goto(url, wait_until='domcontentloaded')
        print("✅ 页面已加载")
        
        # 等待用户登录
        print("\n" + "=" * 60)
        print("⚠️  请在浏览器中完成登录操作")
        print(f"⏱️  程序将等待 {wait_time} 秒...")
        print("💡 登录完成后，程序会自动保存 Cookie")
        print("🔴 按 Ctrl+C 可以提前结束等待")
        print("=" * 60)
        
        try:
            await asyncio.sleep(wait_time)
        except KeyboardInterrupt:
            print("\n\n⚠️  收到中断信号，准备保存 Cookie...")
        
        # 保存 Cookie（storage_state 格式）
        print("\n正在保存 Cookie...")
        cookies = await context.cookies()
        
        # 延长 cookie 过期时间
        current_time = time.time()
        extended_expiry = current_time + (extend_expiry_days * 24 * 60 * 60)
        
        # 验证和清理 cookie 数据
        valid_cookies = []
        for cookie in cookies:
            # 确保必需字段存在
            if 'name' in cookie and 'value' in cookie and 'domain' in cookie:
                # 延长过期时间
                if 'expires' in cookie and cookie['expires'] != -1:
                    cookie['expires'] = extended_expiry
                
                # 确保 sameSite 字段有效（Playwright 要求）
                if 'sameSite' in cookie:
                    if cookie['sameSite'] not in ['Strict', 'Lax', 'None']:
                        cookie['sameSite'] = 'Lax'
                
                valid_cookies.append(cookie)
        
        # 创建 Playwright storage_state 格式
        storage_state = {
            'cookies': valid_cookies,
            'origins': []
        }
        
        cookie_path = Path(cookie_file)
        with open(cookie_path, 'w', encoding='utf-8') as f:
            json.dump(storage_state, f, ensure_ascii=False, indent=2)
        
        print(f"✅ Storage State 已保存!")
        print(f"📁 文件位置: {cookie_path.absolute()}")
        print(f"📊 Cookie 数量: {len(valid_cookies)} 条")
        print(f"⏰ Cookie 过期时间已延长至: {extend_expiry_days} 天后")
        print(f"💡 格式: Playwright storage_state（可直接用于 Browser Use）")
        
        # 显示部分 Cookie 信息（用于验证）
        print("\n📋 Cookie 预览（前 5 条）:")
        for i, cookie in enumerate(cookies[:5], 1):
            print(f"  {i}. {cookie.get('name', 'N/A')}: {cookie.get('domain', 'N/A')}")
        
        if len(cookies) > 5:
            print(f"  ... 还有 {len(cookies) - 5} 条")
        
        print("\n" + "=" * 60)
        print("✅ Cookie 获取完成!")
        print("=" * 60)
        
    except Exception as e:
        print(f"\n❌ 错误: {e}")
        import traceback
        traceback.print_exc()
        
    finally:
        if context:
            try:
                await context.close()
            except:
                pass
        if browser:
            print("\n正在关闭浏览器...")
            try:
                await browser.close()
                print("✅ 浏览器已关闭")
            except:
                pass


async def load_cookies_from_file(cookie_file: str = "douyin_storage_state.json"):
    """
    从文件加载 Storage State（仅读取，不需要浏览器）
    
    参数:
        cookie_file: Storage State 文件路径
        
    返回:
        Storage State 字典，如果文件不存在则返回 None
    """
    try:
        cookie_path = Path(cookie_file)
        
        if not cookie_path.exists():
            print(f"⚠️  Storage State 文件不存在: {cookie_path.absolute()}")
            return None
        
        with open(cookie_path, 'r', encoding='utf-8') as f:
            storage_state = json.load(f)
        
        cookies = storage_state.get('cookies', [])
        print(f"✅ 已读取 Storage State 文件: {cookie_path.absolute()}")
        print(f"📊 Cookie 数量: {len(cookies)} 条")
        
        return storage_state
        
    except Exception as e:
        print(f"❌ 读取 Storage State 文件失败: {e}")
        return None


def main():
    """
    命令行入口
    """
    import argparse
    
    parser = argparse.ArgumentParser(description="Cookie 管理工具")
    parser.add_argument(
        "--url",
        default="https://www.douyin.com",
        help="要访问的网址（默认: https://www.douyin.com）"
    )
    parser.add_argument(
        "--output",
        default="douyin_storage_state.json",
        help="Cookie 保存的文件名（默认: douyin_storage_state.json）"
    )
    parser.add_argument(
        "--wait",
        type=int,
        default=60,
        help="等待用户登录的时间（秒，默认: 60）"
    )
    parser.add_argument(
        "--extend-days",
        type=int,
        default=365,
        help="延长 cookie 过期时间的天数（默认: 365）"
    )
    
    args = parser.parse_args()
    
    # 运行 Cookie 获取
    asyncio.run(get_and_save_cookies(
        url=args.url,
        cookie_file=args.output,
        wait_time=args.wait,
        extend_expiry_days=args.extend_days
    ))


if __name__ == "__main__":
    main()
